import os
import logging
import pandas as pd
from langchain_community.utilities.sql_database import SQLDatabase

class PostgreSQLConnector:
    """PostgreSQL数据库连接器
    负责创建数据库连接、获取表结构信息等操作
    """
    def __init__(self):
        self.db_uri = os.getenv("DB_URI")
        self.allowed_tables = os.getenv("ALLOWED_TABLES", "").split(",")
        self.db = self._create_connection()
        
    def _create_connection(self):
        """创建数据库连接
        
        Returns:
            SQLDatabase: 数据库连接实例
        
        Raises:
            Exception: 连接失败时抛出异常
        """
        try:
            return SQLDatabase.from_uri(
                self.db_uri,
                include_tables=self.allowed_tables if self.allowed_tables[0] else None,
                sample_rows_in_table_info=0  # 不获取数据样本，保护隐私
            )
        except Exception as e:
            logging.error(f"数据库连接失败: {str(e)}")
            raise
    
    def get_table_info(self):
        """获取表结构信息
        
        Returns:
            str: 表结构信息字符串
        """
        return self.db.get_table_info()
    
    def get_connection(self):
        """获取数据库连接实例
        
        Returns:
            SQLDatabase: 数据库连接实例
        """
        return self.db

    def exp_data_scv(self,query,file_name) -> str:
        """导出查询数据到CSV文件
        Args:
            query (str): SQL查询语句
            file_name (str): 输出文件名
        Returns:
            filepath: 导出文件路径
        """
        if not query:
            logging.error("exp_data_scv,查询语句为空")
            return  "导出失败,查询语句为空"
        if not file_name:
            logging.error("exp_data_scv,文件名为空")
            return "导出失败,文件名为空"    
        # 检查文件扩展名
        if not file_name.endswith('.csv'):
            file_name += '.csv'

        # 导出到CSV
        export_dir = os.getenv("EXPORT_DIR")
        os.makedirs(export_dir, exist_ok=True)
        logging.info(f"exp_data_scv 导出目录: {export_dir}")
        file_path = os.path.join(export_dir, file_name)
        logging.info(f"exp_data_scv 导出文件路径: {file_path}")

        try:
            # 执行查询
            result = self.db.run(query)
            
            # 检查结果类型并处理
            if isinstance(result, str):
                # 如果结果是字符串，尝试将其解析为DataFrame
                # 这里需要根据实际返回的字符串格式进行解析
                # 假设结果是JSON格式或类似表格的数据
                logging.info(f"exp_data_scv 查询结果是字符串，尝试解析为DataFrame。")
                try:
                    # 尝试使用pandas的read_sql_query直接执行查询
                    # 注意：这里需要获取底层连接对象，具体实现可能需要调整
                    from sqlalchemy import create_engine
                    engine = create_engine(self.db_uri)
                    df = pd.read_sql_query(query, engine)
                except Exception as e:
                    logging.error(f"exp_data_scv 使用pandas读取SQL失败: {str(e)}")
                    # 尝试其他解析方式
                    # 这里仅作为示例，实际解析逻辑需要根据实际情况调整
                    # 假设有更复杂的解析逻辑
                    raise ValueError(f"exp_data_scv 无法将查询结果转换为DataFrame: {str(e)}")
            elif hasattr(result, 'to_csv'):
                # 如果结果已经有to_csv方法，直接使用
                logging.info(f"exp_data_scv 查询结果有to_csv方法")
                df = result
            else:
                # 其他情况，尝试转换为DataFrame
                logging.info(f"exp_data_scv 查询结果类型: {type(result)}，尝试转换为DataFrame。")
                df = pd.DataFrame(result)
                
            # 导出为CSV
            df.to_csv(file_path, index=False)
            logging.info(f"exp_data_scv 数据已导出到 {file_path},数据量: {len(df)}")
            return file_path
        except Exception as e:
            logging.error(f"exp_data_scv 导出失败: {str(e)}")
            return f"导出失败,{str(e)}"
